package com.example.spider.provider.m3u8.handler;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.StrUtil;
import com.example.spider.cluster.redis.RedisBatchProgressCache;
import com.example.spider.config.ThreadPoolConfiguration;
import com.example.spider.provider.m3u8.M3u8FileNameTemplate;
import com.example.spider.provider.m3u8.M3u8Parser;
import com.example.spider.provider.m3u8.dto.M3u8Key;
import com.example.spider.provider.m3u8.dto.M3u8MetaInfo;
import com.example.spider.provider.m3u8.dto.M3u8TsPart;
import com.example.spider.provider.m3u8.task.M3u8KeyTask;
import com.example.spider.provider.m3u8.task.M3u8MetaInfoParserTask;
import com.example.spider.provider.m3u8.task.M3u8TsTask;
import com.example.spider.task.SpiderResultHandler;
import com.example.spider.task.SpiderTaskQueue;
import com.example.spider.task.Task;
import com.example.spider.task.limit.FixWindowLimit;
import com.example.spider.util.PathUtil;
import com.example.spider.web.ProgressController;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.lang3.time.DurationFormatUtils;
import org.shoulder.batch.progress.BatchProgress;
import org.shoulder.batch.progress.BatchProgressCache;
import org.shoulder.batch.progress.Progress;
import org.shoulder.core.concurrent.Threads;
import org.shoulder.core.exception.CommonErrorCodeEnum;
import org.shoulder.core.util.AssertUtils;
import org.shoulder.core.util.JsonUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardCopyOption;
import java.nio.file.StandardOpenOption;
import java.time.Duration;
import java.util.BitSet;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Semaphore;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * 处理 m3u8 内容，初始化密钥，生产TS任务，合并任务
 *
 * @author lym
 */
@Slf4j
@Component
public class M3u8FirstHandler implements SpiderResultHandler<M3u8MetaInfoParserTask, String> {

    static {
        Task.registerType(M3u8MetaInfoParserTask.class);
    }

    @Qualifier(ThreadPoolConfiguration.HANDLER)
    @Autowired
    private ThreadPoolExecutor executor;

    @Value("${spider.m3u8.alwaysReMergeTsPart:false}")
    private boolean alwaysReMergeTsPart;

    @Autowired
    private BatchProgressCache progressCache;

    @Autowired
    private M3u8FileNameTemplate m3u8FileNameTemplate;

    // 用于限制磁盘的使用，这里设置为3，仅供参考
    private static final Semaphore semaphore = new Semaphore(3);

    private boolean alwaysDownloadKey = true;

    private static CountDownLatch completeKey(M3u8MetaInfoParserTask task, M3u8MetaInfo metaInfo) {
        if (MapUtils.isEmpty(metaInfo.getKeyMap())) {
            return new CountDownLatch(0);
        }
        CountDownLatch latch = new CountDownLatch(metaInfo.getKeyMap().size());
        // 大部分m3u8文件只有一个key
        for (M3u8Key m3u8Key : metaInfo.getKeyMap().values()) {
            if (!m3u8Key.isInit()) {
                // 需要初始化的
                if (m3u8Key.isEncrypted()) {
                    SpiderTaskQueue.putHighestPriorityTask(new M3u8KeyTask()
                            .setKey(m3u8Key)
                            .setLatch(latch)
                            .setUrl(m3u8Key.getUrl())
                            .setParent(task)
                            .setLimitStrategy(task.getLimitStrategy() == null ? null : task.getLimitStrategy().copyWithNewKey(task.getLimitStrategy().getKey() + "_key"))
                    );
                    continue;
                }

                // 不需要初始化的
                m3u8Key.setInit(true);
                latch.countDown();
            }
        }

        return latch;
    }

    /**
     * 考虑到合并过程比较长（秒级），该方法会保证产出的目标文件不会缺失内容，即是合并时出现宕机，下次执行会自动删除错误文件，重新合并
     * 一旦产出目标文件，一定是内容完整的
     */
    private static void mergeTsParts(String tsPartsDir, Path resultTsFilePath) {
        File[] tsPartFileArr = new File(tsPartsDir).listFiles();
        // 刚下载的肯定不为空
        assert tsPartFileArr != null;

        log.debug("merge tsFile start");
        Path tempResultTsFilePath = Path.of(resultTsFilePath.toString() + File.separator + ".tmp");
        try {
            Files.deleteIfExists(tempResultTsFilePath);
            Files.deleteIfExists(resultTsFilePath);
            PathUtil.ensureExist(tempResultTsFilePath.getParent());
            Files.createFile(tempResultTsFilePath);
            long start = System.currentTimeMillis();
            for (File tsFile : tsPartFileArr) {
                // 清理临时文件
                if (tsFile.getName().endsWith(".tmp")) {
                    tsFile.delete();
                    continue;
                }
                // 读 ts 分片，写到临时文件
                Path readPath = Path.of(tsFile.getAbsolutePath());
                Files.write(tempResultTsFilePath, Files.readAllBytes(readPath), StandardOpenOption.APPEND);
            }
            // 写完搬迁到正式文件(原子)
            Files.move(tempResultTsFilePath, resultTsFilePath, StandardCopyOption.ATOMIC_MOVE, StandardCopyOption.REPLACE_EXISTING, StandardCopyOption.COPY_ATTRIBUTES);

            long cost = System.currentTimeMillis() - start;
            log.info("merge finished: {} cost {}ms.", resultTsFilePath, cost);
        } catch (Exception e) {
            log.error("merge tsFile ex", e);
        }
    }

    private static void createTsPartTasks(M3u8MetaInfoParserTask task, String partTsFileNamePrefix, List<M3u8TsPart> tsPart, Progress progress) {
        Path dirPath = Path.of(partTsFileNamePrefix);
        PathUtil.ensureExist(dirPath);

        int partSize = tsPart.size();
        double totalSecond = 0;
        // 每个ts part + 最后merged是否完成的标志
        task.setTsPartDownloadStatus(new BitSet(partSize + 1));
//        CountDownLatch tsDownloadLatch = new CountDownLatch(partSize);
        int maxLength = String.valueOf(partSize).length();
        for (int i = 0; i < partSize; i++) {
            String index = String.valueOf(i);
            int currentLength = index.length();

            StringBuilder fileName = new StringBuilder();
            while (fileName.length() < maxLength - currentLength) {
                fileName.append("0");
            }
            fileName.append(i);
            fileName.append(".ts");
            String tsFilePath = partTsFileNamePrefix + fileName;

            M3u8TsPart part = tsPart.get(i);
            totalSecond += Double.parseDouble(part.getDuration());
            part.setFilePath(tsFilePath);
            if (FileUtil.exist(tsFilePath) && FileUtil.size(new File(tsFilePath)) > 100L) {
                boolean finishedInHistory = task.getTsPartDownloadStatus().get(part.getPartIndex());
//                task.getTsPartDownloadStatus().set(part.getPartIndex());
                // 去重，可重跑，支持极端场景幂等：创建了ts、未写入采用 finishedInHistory，覆盖原有tsPart文件
                log.debug("already exists, maybe had downloaded, reuse file: {}", tsFilePath);
//                tsDownloadLatch.countDown();
                progress.finishPart(i);
                continue;
            }
//            ((FixWindowLimit) task.getLimitStrategy()).setLimitDuration(Duration.ofMillis(20));
            M3u8TsTask tsTask = new M3u8TsTask()
                    .setTsPart(part)
                    .setProgressId(progress.getId())
//                    .setCountDownLatch(tsDownloadLatch)
//                    .setTsPartDownloadStatus(task.getTsPartDownloadStatus())
//                    .setLimitStrategy(task.getLimitStrategy() == null ? null : task.getLimitStrategy().copyWithNewKey(task.getLimitStrategy().getKey() + "_ts"))
                    .setLimitStrategy(new FixWindowLimit("", Duration.ofMillis(100)))
                    .setUrl(part.getUrl())
                    .setIdempotentFilePath(tsFilePath);
            //crossSiteHeaders.forEach(tsTask::addHeader);
            SpiderTaskQueue.putTask(tsTask);
        }
        if (totalSecond < 30) {
            // 时长太小，仅打印warn，有需要可以个性化定制过短/过长的视频
            log.warn("[unexpected duration] {}s, tsSize={}, ", ((long) totalSecond), partSize);
        }
        log.debug("finished create ts task.");
    }

    /**
     * 分析列表结果，创建下载详情任务
     */
    @Override
    public void handle(M3u8MetaInfoParserTask task, String result) throws IOException {
        // https://api.zhiyinapi.com/api/v1/get-keys
        // https://cdn77-vid.xvideos-cdn.com/6INcvsbeaRA6Vb8ixqdZQg==,1677525576/videos/hls/14/f2/11/14f2117ad464e4e2272e136923083132/hls-480p-9e63e.m3u8
        log.debug("prepare handing m3u8 meta " + task.getUrl());

        long startTime = System.currentTimeMillis();

        // 解析
        M3u8MetaInfo cachedMetaInfo = task.getCachedM3u8MetaInfo();
        boolean alreadyDownloadKey = cachedMetaInfo != null;
        M3u8MetaInfo metaInfo = alreadyDownloadKey ? cachedMetaInfo : M3u8Parser.parse(result, task.getUrl(), task.getForbiddenPartUrls());

        long parseCost = System.currentTimeMillis() - startTime;
        log.debug("parseM3u8Cost={}ms.", parseCost);

        // subM3u8
        if (CollectionUtil.isNotEmpty(metaInfo.getSubMetaInfo())) {
            // 说明本次请求只是菜单，非ts内容
            AssertUtils.empty(metaInfo.getTsPart(), CommonErrorCodeEnum.ILLEGAL_PARAM);
            // todo P2 选一个清晰度下载还是都下载？
            for (int i = 0; i < metaInfo.getSubMetaInfo().size(); i++) {
                M3u8MetaInfo subInfo = metaInfo.getSubMetaInfo().get(i);
                String subId = task.getId() + "_" + i;
                M3u8MetaInfoParserTask m3u8MetaInfoParserTask = new M3u8MetaInfoParserTask()
                        .setId(subId)
                        .setMetaInfo(metaInfo)
                        .setForbiddenPartUrls(task.getForbiddenPartUrls())
                        .setCrossSiteHeaders(task.getCrossSiteHeaders())
                        .setSourcePageUrl(task.getSourcePageUrl())
                        .setLimitStrategy(task.getLimitStrategy())
                        .setUrl(subInfo.getUrl());
                SpiderTaskQueue.putTask(m3u8MetaInfoParserTask);
                // 默认只下载第一个
                if (metaInfo.getSubMetaInfo().size() > 1) {
                    log.warn("m3u8 {}个子链接，只下载第一个.", metaInfo.getSubMetaInfo().size());
                }
                return;
            }
        }

        String m3u8FileName = m3u8FileNameTemplate.calculateM3u8FileName(task.getId());
        String metaFileName = m3u8FileNameTemplate.calculateMetaFileName(task.getId());
        String partTsFileDir = m3u8FileNameTemplate.calculateTsPartFileDir(task.getId());
        String videoFileName = m3u8FileNameTemplate.calculateVideoFileName(task.getId());
        Path taskSavePath = Path.of(m3u8FileNameTemplate.calculateTaskFileName(task.getId()));

        // 初始化密钥
        if (!alwaysDownloadKey && alreadyDownloadKey) {
            log.info("SKIP downloadKey: alreadyDownload.");
        } else {
            CountDownLatch keyLatch = completeKey(task, metaInfo);

            // 等待下载的同时保存m3u8任务与文件，便于下次断点续传 - 保存m3u8内容
            // task 第一次写入
            PathUtil.rewriteFile(JsonUtils.toJson(task, ObjectMapper::writerWithDefaultPrettyPrinter), taskSavePath);
            PathUtil.rewriteFile(result, Path.of(m3u8FileName));

            // 等待 key 加载完成
            boolean finished = waitTillFinished(keyLatch, 5, TimeUnit.MINUTES) && (MapUtils.isEmpty(metaInfo.getKeyMap()) || metaInfo.getKeyMap().values().stream().allMatch(M3u8Key::isInit));
            // 不重试，宣布失败
            if (!finished) {
                if (task.getRecoverCount() < 100) {
                    task.setRecoverCount(99999);
                }
                task.setRecoverCount(task.getRecoverCount() + 1);
            }
            // 等待重试
            if (!finished) {
                log.warn("RETRY({} - {})- waiting download m3u8Key timeout, ", task.getClass().getSimpleName(), task.getUrl());
                Threads.delay(() -> SpiderTaskQueue.putTask(task), Duration.ofMillis(5));
                return;
            }
        }

        // 创建爬TS的任务 此处允许重复创建，后续会幂等
        Progress progress = progressCache.getProgress(task.getId());
        if(!(progressCache instanceof RedisBatchProgressCache) && progress == null) {
            // 内存缓存
            BatchProgress memoryProgress = new BatchProgress();
            memoryProgress.setId(task.getId());
            progress = memoryProgress;
        }

        progress.setTotal(metaInfo.getTsPart().size());

        startTime = System.currentTimeMillis();
        createTsPartTasks(task, partTsFileDir, metaInfo.getTsPart(), progress);
        long createTsPartTasksCost = System.currentTimeMillis() - startTime;
        log.debug("createTsPartTasksCost={}.", createTsPartTasksCost);
        final long downloadStartTime = startTime;
        progress.setOnFinishCallback((progressId, p) -> mergeTsAfterDownload(task, downloadStartTime, metaInfo));
        progress.start();
        progressCache.triggerFlushProgress(progress);
        if (!progress.hasFinish()) {
            log.warn("m3u8({}) 开始下载，进度： {} ", metaInfo.getUrl(), ProgressController.getProgressUrl(task.getId()));
            // 保存触发这个m3u8的任务与解析结果（来源，m3u8 url）便于下次断点续传
            PathUtil.rewriteFile(JsonUtils.toJson(metaInfo, ObjectMapper::writerWithDefaultPrettyPrinter), Path.of(metaFileName));
        }

        log.info("Downloading... m3u8({}) parseM3u8Cost={}ms. createTsPartTasksCost={}ms. file={}",
                task.getId(), parseCost, createTsPartTasksCost, videoFileName);
    }

    private void mergeTsAfterDownload(M3u8MetaInfoParserTask task, long startTime, M3u8MetaInfo metaInfo) {
        String partTsFileDir = m3u8FileNameTemplate.calculateTsPartFileDir(task.getId());
        String videoFileName = m3u8FileNameTemplate.calculateVideoFileName(task.getId());
        Path taskSavePath = Path.of(m3u8FileNameTemplate.calculateTaskFileName(task.getId()));

        // task 更新状态
        PathUtil.rewriteFile(JsonUtils.toJson(task, ObjectMapper::writerWithDefaultPrettyPrinter), taskSavePath);
        long tsDownloadCost = (System.currentTimeMillis() - startTime);
        String downloadTsCostFormatted = DurationFormatUtils.formatDurationWords(tsDownloadCost, true, true);

        // 合并 TS
        mergeTsToResult(task, metaInfo, partTsFileDir, videoFileName, taskSavePath);

        // 创建空内容，文件名=片名的文件
        if (StrUtil.isNotBlank(task.getName())) {
            PathUtil.rewriteFile("[just a name mark file]\n" + task.getName(), taskSavePath.getParent().resolve(task.getName()));
        }

        log.info("FINISHED download m3u8({})  downloadCost={}. mergeTo file={}", task.getId(), downloadTsCostFormatted, videoFileName);
    }

    private void mergeTsToResult(M3u8MetaInfoParserTask task, M3u8MetaInfo metaInfo, String partTsFileDir, String videoFileName,
                                              Path taskSavePath) {
        Path mergeResultFilePath = Path.of(videoFileName);
        boolean mergedInHistory = Files.exists(mergeResultFilePath) && task.getTsPartDownloadStatus().get(metaInfo.getTsPart().size() + 1);
        if (!mergedInHistory || alwaysReMergeTsPart) {
            // 建议限制到这一步的个数，否则顺序访问变为随机访问性能下降
            try {
                semaphore.acquire();
                mergeTsParts(partTsFileDir, mergeResultFilePath);

                task.getTsPartDownloadStatus().set(metaInfo.getTsPart().size() + 1);
                // task 第三次写入覆盖状态
                PathUtil.rewriteFile(JsonUtils.toJson(task, ObjectMapper::writerWithDefaultPrettyPrinter), taskSavePath);
                // feature 合并后删除ts？暂时不支持，必要性不大，后续提供扫描删除脚本即可
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            } finally {
                semaphore.release();
            }

        } else {
            log.info("SKIP merge: already merged in history.");
        }
    }

    /**
     * todo P2 feat 返回：失败原因、可否重试
     */
    private boolean waitTillFinished(CountDownLatch latch, long timeout, TimeUnit unit) {
        synchronized (executor) {
            executor.setMaximumPoolSize(executor.getMaximumPoolSize() + 1);
            executor.setCorePoolSize(executor.getCorePoolSize() + 1);
        }
        try {
            return latch.await(timeout, unit);
        } catch (InterruptedException e) {
            return false;
        } finally {
            synchronized (executor) {
                executor.setMaximumPoolSize(executor.getMaximumPoolSize() - 1);
                executor.setCorePoolSize(executor.getCorePoolSize() - 1);
            }
        }
    }
}
